"""TensorRT核心功能模块"""
import os
import json
import tensorrt as trt
from typing import Optional, Dict, Any
from abc import ABC, abstractmethod

class TensorRTLogger:
    """TensorRT日志管理器"""
    _instance = None
    
    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
            cls._instance.logger = trt.Logger(trt.Logger.INFO)
        return cls._instance
    
    @property
    def trt_logger(self):
        return self.logger

class TensorRTBuilder:
    """TensorRT引擎构建器基类"""
    
    def __init__(self, workspace_size: int = 1 << 32):
        self.logger = TensorRTLogger().trt_logger
        self.workspace_size = workspace_size
        self.builder = None
        self.network = None
        self.config = None
        
    def _create_network(self) -> trt.INetworkDefinition:
        """创建网络定义"""
        if self.builder is None:
            self.builder = trt.Builder(self.logger)
        
        network_flags = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
        self.network = self.builder.create_network(network_flags)
        return self.network
    
    def _create_config(self) -> trt.IBuilderConfig:
        """创建构建配置"""
        self.config = self.builder.create_builder_config()
        self.config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, self.workspace_size)
        return self.config
    
    def _parse_onnx(self, onnx_path: str) -> bool:
        """解析ONNX模型"""
        if not os.path.exists(onnx_path):
            raise FileNotFoundError(f"ONNX file not found: {onnx_path}")
        
        parser = trt.OnnxParser(self.network, self.logger)
        
        with open(onnx_path, 'rb') as f:
            if not parser.parse(f.read()):
                error_msgs = []
                for idx in range(parser.num_errors):
                    error_msgs.append(str(parser.get_error(idx)))
                raise RuntimeError(f"ONNX parsing failed: {'; '.join(error_msgs)}")
        
        return True
    
    def _build_and_serialize(self) -> bytes:
        """构建并序列化引擎"""
        serialized_engine = self.builder.build_serialized_network(self.network, self.config)
        if serialized_engine is None:
            raise RuntimeError("Failed to build serialized engine")
        return serialized_engine
    
    def save_engine(self, serialized_engine: bytes, engine_path: str) -> None:
        """保存引擎到文件"""
        os.makedirs(os.path.dirname(engine_path), exist_ok=True)
        with open(engine_path, 'wb') as f:
            f.write(serialized_engine)
        print(f"Engine successfully saved to: {engine_path}")

class QuantizationStrategy(ABC):
    """量化策略抽象基类"""
    
    @abstractmethod
    def apply_quantization(self, config: trt.IBuilderConfig, network: trt.INetworkDefinition) -> None:
        """应用量化策略"""
        pass

class FP16Strategy(QuantizationStrategy):
    """FP16量化策略"""
    
    def apply_quantization(self, config: trt.IBuilderConfig, network: trt.INetworkDefinition) -> None:
        config.set_flag(trt.BuilderFlag.FP16)

class INT8Strategy(QuantizationStrategy):
    """INT8量化策略"""
    
    def __init__(self, calibrator: Optional[trt.IInt8Calibrator] = None):
        self.calibrator = calibrator
    
    def apply_quantization(self, config: trt.IBuilderConfig, network: trt.INetworkDefinition) -> None:
        config.set_flag(trt.BuilderFlag.INT8)
        if self.calibrator:
            config.int8_calibrator = self.calibrator